#include <windows.h>
#include <tlhelp32.h>

#include "memwtr.h"

HANDLE phandle = NULL;

/*
 * Class:     conch_yaoms_jni_MemoryWatcher
 * Method:    isReady
 * Signature: ()Z
 */
JNIEXPORT jboolean JNICALL Java_conch_yaoms_jni_MemoryWatcher_isReady
    (JNIEnv * env, jobject obj) {
        if (phandle != NULL && phandle != INVALID_HANDLE_VALUE)
            return 1;
        else
            return 0;
}

/*
 * Class:     conch_yaoms_jni_MemoryWatcher
 * Method:    openProcess
 * Signature: (I)V
 */
JNIEXPORT void JNICALL Java_conch_yaoms_jni_MemoryWatcher_openProcess
(JNIEnv * env, jobject obj, jint process_id) {

    if (!Java_conch_yaoms_jni_MemoryWatcher_isReady(env, obj))
        phandle = OpenProcess(PROCESS_VM_READ,FALSE,process_id);
    //phandle = OpenProcess(PROCESS_ALL_ACCESS,FALSE,process_id);
}

/*
 * Class:     conch_yaoms_jni_MemoryWatcher
 * Method:    closeProcess
 * Signature: ()V
 */
JNIEXPORT void JNICALL Java_conch_yaoms_jni_MemoryWatcher_closeProcess
(JNIEnv * env, jobject obj) {

    if (Java_conch_yaoms_jni_MemoryWatcher_isReady(env, obj)) {
        // Close HANDLE
        CloseHandle(phandle);
        phandle = NULL;
    }
}


/*
 * Class:     conch_yaoms_jni_MemoryWatcher
 * Method:    getProcessList
 * Signature: ()Ljava/util/ArrayList;
 */
JNIEXPORT jobject JNICALL Java_conch_yaoms_jni_MemoryWatcher_getProcessList
(JNIEnv * env, jobject obj) {

    jclass list_clz = env->FindClass("java/util/ArrayList");
    jmethodID list_init = env->GetMethodID(list_clz, "<init>", "()V");
    jmethodID list_add = env->GetMethodID(list_clz, "add","(Ljava/lang/Object;)Z");

    jclass process_item_clz = env->FindClass("conch/yaoms/jni/ProcessItem");
    jmethodID item_init = env->GetMethodID(process_item_clz, "<init>", "()V");
    jfieldID process_item_name = env->GetFieldID(process_item_clz, "processImageName", "Ljava/lang/String;");
    jfieldID process_item_id = env->GetFieldID(process_item_clz, "processId", "I");

    jobject list_obj = env->NewObject(list_clz,list_init);

    PROCESSENTRY32 pe32 = {sizeof (PROCESSENTRY32)};

    HANDLE hProcSnap = CreateToolhelp32Snapshot (TH32CS_SNAPPROCESS, 0);
    if (hProcSnap == INVALID_HANDLE_VALUE) return 0;

    if (Process32First (hProcSnap, &pe32))
        do {
            jobject item = env->NewObject(process_item_clz, item_init);

            jclass str_clz = env->FindClass("Ljava/lang/String;");
            jmethodID str_init = env->GetMethodID(str_clz, "<init>", "([BLjava/lang/String;)V");

            jbyteArray bytes = env->NewByteArray(strlen(pe32.szExeFile));
            env->SetByteArrayRegion(bytes, 0, strlen(pe32.szExeFile), (jbyte*)pe32.szExeFile);

            jstring encoding = env->NewStringUTF("utf-8");
            jstring name = (jstring) env->NewObject(str_clz, str_init, bytes, encoding);

            env->SetObjectField(item,process_item_name,name);
            env->SetIntField(item,process_item_id, pe32.th32ProcessID);

            env->CallVoidMethod(list_obj,list_add,item);
        } while (Process32Next (hProcSnap, &pe32));

    CloseHandle (hProcSnap);

    return list_obj;
}

/*
 * Class:     conch_yaoms_jni_MemoryWatcher
 * Method:    getBytes
 * Signature: (III)[B
 */
JNIEXPORT jbyteArray JNICALL Java_conch_yaoms_jni_MemoryWatcher_getBytes
(JNIEnv * env, jobject obj, jint process_id, jint address, jint count) {

    jbyteArray bytes = env->NewByteArray(count);

    Java_conch_yaoms_jni_MemoryWatcher_openProcess(env, obj, process_id);

        // Read
        if (!Java_conch_yaoms_jni_MemoryWatcher_isReady(env, obj)) return NULL;

    ReadProcessMemory(phandle,(LPCVOID)address,&bytes,count,0);

    return bytes;
}

